package com.linyaonan.leetcode.easy._637;

import java.util.*;

/**
 * 给定一个非空二叉树, 返回一个由每层节点平均值组成的数组.
 * <p>
 * 示例 1:
 * <p>
 * 输入:
 * 3
 * / \
 * 9  20
 * /  \
 * 15   7
 * 输出: [3, 14.5, 11]
 * 解释:
 * 第0层的平均值是 3,  第1层是 14.5, 第2层是 11. 因此返回 [3, 14.5, 11].
 * 注意：
 * <p>
 * 节点值的范围在32位有符号整数范围内。
 *
 * @author: Lin
 * @date: 2020/1/4
 */
public class AverageOfLevels {
    Map<Integer, Double> map = new HashMap<>();
    Map<Integer, Integer> countMap = new HashMap<>();

    public List<Double> averageOfLevels(TreeNode root) {
        bfs(root, 1);
        int size = map.entrySet().size();
        List<Double> result = new ArrayList<>();
        for (int i = 1; i <= size; i++) {
            Double sum = map.get(i);
            result.add(sum / countMap.get(i));
        }
        return result;
    }

    private void bfs(TreeNode node, int count) {
        if (node != null) {
            map.put(count, map.getOrDefault(count, 0D) + node.val);
            countMap.put(count, countMap.getOrDefault(count, 0) + 1);
            count++;
            bfs(node.left, count);
            bfs(node.right, count);
        }
    }

    public List<Double> averageOfLevels2(TreeNode root) {
        List<Double> res = new ArrayList<>();
        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(root);
        while (!queue.isEmpty()) {
            long sum = 0, count = 0;
            Queue<TreeNode> temp = new LinkedList<>();
            while (!queue.isEmpty()) {
                TreeNode n = queue.remove();
                sum += n.val;
                count++;
                if (n.left != null) {
                    temp.add(n.left);
                }
                if (n.right != null) {
                    temp.add(n.right);
                }
            }
            queue = temp;
            res.add(sum * 1.0 / count);
        }
        return res;
    }
}
